library(Seurat)
library(velocyto.R)
library(SeuratWrappers)
source('00_util_scripts/mod_seurat.R')
library(tidyverse)
library(tidyseurat)

# seekgene starsolo velocyto ----------
skg_path <- '~/work/results/seekgene_crc/expr_mat/'

skg_barcode <- list.files(skg_path, pattern = '^barcodes',full.names = TRUE,recursive = TRUE) |>
  str_subset('Velocyto/filtered/bar')

skg_feature <- list.files(skg_path, pattern = '^feature',full.names = TRUE,recursive = TRUE) |>
  str_subset('Velocyto/filtered/fea')

skg_spliced <- list.files(skg_path, pattern = '^spliced',full.names = TRUE,recursive = TRUE) |>
  str_subset('filtered')

skg_unspliced <- list.files(skg_path, pattern = 'unspliced',full.names = TRUE,recursive = TRUE) |>
  str_subset('filtered')

skg_ambi <- list.files(skg_path, pattern = 'ambi',full.names = TRUE,recursive = TRUE) |>
  str_subset('filtered')

skg_all <- list.files(skg_path, pattern = '^matrix',full.names = TRUE,recursive = TRUE) |>
    str_subset('Gene/filtered')

skg_all_barcode <- list.files(skg_path, pattern = '^barcode',
                              full.names = TRUE,
                              recursive = TRUE) |>
  str_subset('Gene/filtered')

skg_all_feature <- list.files(skg_path,
                              pattern = '^feature',
                              full.names = TRUE,
                              recursive = TRUE) |>
  str_subset('Gene/filtered')

mtx_spliced <- list(skg_spliced, skg_barcode, skg_feature) |>
  pmap(ReadMtx, .progress = TRUE) |>
  map2(c('crc221125','crc230304','crc230309'), add_name_field)

mtx_unspliced <- list(skg_unspliced, skg_barcode, skg_feature) |>
  pmap(ReadMtx, .progress = TRUE) |>
  map2(c('crc221125','crc230304','crc230309'), add_name_field)

mtx_all <- list(skg_all, skg_barcode, skg_feature) |>
  pmap(ReadMtx, .progress = TRUE) |>
  map2(c('crc221125','crc230304','crc230309'), add_name_field)

mtx_all_gene <- list(skg_all, skg_all_barcode, skg_all_feature) |>
  pmap(ReadMtx, .progress = TRUE) |>
  map2(c('crc221125','crc230304','crc230309'), add_name_field)

# create inital seurat from normal mtx -----
sobj_all <- mtx_all_gene |>
  map(CreateSeuratObject,
      min.cell = 3, min.feature = 200,
      names.field = 2)

add_timing <- function(sobj, spl_mtx, uns_mtx){
  spl <- CreateAssayObject(spl_mtx)
  uns <- CreateAssayObject(uns_mtx)
  Key(spl) <- 'spliced_'
  Key(uns) <- 'unspliced_'
  sobj@assays$spliced <- spl
  sobj@assays$unspliced <- uns
  sobj
}

sobj <- list(sobj_all, mtx_spliced, mtx_unspliced) |>
  pmap(add_timing) |>
  purrr::reduce(merge)

sobj

sobj$mt.ratio <- PercentageFeatureSet(sobj, 'MT-')

sobj <- sobj |>
  filter(nFeature_RNA >= 200 & mt.ratio < 10)

sobj <- sobj |>
  quick_process_seurat()

sobj <- RunTSNE(sobj, reduction = "harmony", dims = 1:20)

DimPlot(sobj, reduction = 'tsne')

hpca <- celldex::HumanPrimaryCellAtlasData()

sobj <- mark_cell_type_singler(sobj, hpca, new_label = 'hpca_main')

DimPlot(sobj, group.by = 'hpca_main', label = TRUE)

# remove stroma cell and re-cluster ----------
sobj <- sobj |>
  filter(!str_detect(hpca_main, 'Endothe|Epithe|Fibro|stem'))

Idents(sobj) <- 'singler_label'

celltype_markers <- FindAllMarkers(sobj, only.pos = TRUE, logfc.threshold = 1, min.pct = 0.3)

best_markers <- celltype_markers |>
  group_by(cluster) |>
  slice_max(avg_log2FC, n = 2) |>
  pull(gene) |> unique()

sobj |> DotPlot(features = best_markers) +
  theme(axis.text.x = element_text(angle = 45, hjust = 1, vjust = 1))

write_rds(sobj, 'CRC-I/data/seekgene/crc-starsolo-velo.rds')
# when only include normal expr
write_rds(sobj, 'CRC-I/data/seekgene/crc-starsolo.rds')

# re-cluster on immune cells -------------
sobj <- read_rds('CRC-I/data/seekgene/crc-starsolo.rds')

sobj <- sobj |>
  quick_process_seurat()

monaco <- celldex::MonacoImmuneData()

sobj <- mark_cell_type_singler(sobj, monaco, fine_label = TRUE, new_label = 'monaco_fine')

sobj <- RunTSNE(sobj, reduction = "harmony", dims = 1:20)

DimPlot(sobj, reduction = 'tsne', group.by = 'monaco_fine')

FeaturePlot(sobj, features = c('CD4','CD8A'))

cell.dist <- as.dist(1-armaCor(t(sobj@reductions$pca@cell.embeddings)))

# set ncores to speed up!
# @1 = 299s @8 = 52s @16 = 42s @32 = 29s
system.time(sobj <- sobj |> RunVelocity(cell.dist = cell.dist, ncores = 32))

# DO NOT set ncores
# @1 = 66s @2 = 226s
show.velocity.on.embedding.cor(emb = Embeddings(sobj, reduction = "umap"),
                               vel = Tool(sobj, slot = "RunVelocity"),
                               n = 200,
                               scale = "sqrt", 
                               cex = 0.8,
                               arrow.scale = 3,
                               show.grid.flow = TRUE,
                               min.grid.cell.mass = 0.5,
                               grid.n = 40,
                               do.par = FALSE,
                               cell.border.alpha = 0.1) |> system.time()
